//! Type context information.

use crate::{EnvExpr, LiteralValue};

use super::ValueType;

/// Argument to some SSA expression.
#[derive(Clone, Debug)]
pub enum SsaArg {
    /// A literal value.
    Literal(LiteralValue),

    /// A previously assigned local variable.
    Local(LocalRef),

    /// An argument to the function.
    Arg(ArgRef),

    /// A reference to a field.
    Field(FieldRef),
}

/// General type context across compilation.
pub trait GenericContext {
    // TODO
}

/// Type context local to a specific module.
pub trait ModuleContext: GenericContext {
    // TODO
}

/// Field references local to some container.
///
/// Note that these are real fields in the native representation!  These are not
/// encoded fields that we extract from within a buffer.
#[derive(Copy, Clone, Debug, Hash, Eq, PartialEq, Ord, PartialOrd)]
pub struct FieldRef(u16);

impl FieldRef {
    pub fn inner(&self) -> u16 {
        self.0
    }
}

impl From<u16> for FieldRef {
    fn from(value: u16) -> Self {
        Self(value)
    }
}

impl From<usize> for FieldRef {
    fn from(value: usize) -> Self {
        if value > u16::MAX as usize {
            panic!("tysys: out of bounds field ref ({value})");
        }
        Self::from(value as u16)
    }
}

/// Argument reference local to some function.
#[derive(Copy, Clone, Debug, Hash, Eq, PartialEq, Ord, PartialOrd)]
pub struct ArgRef(u16);

impl ArgRef {
    pub fn inner(&self) -> u16 {
        self.0
    }
}

impl From<u16> for ArgRef {
    fn from(value: u16) -> Self {
        Self(value)
    }
}

/// Local variable in some function.
#[derive(Copy, Clone, Debug, Hash, Eq, PartialEq, Ord, PartialOrd)]
pub struct LocalRef(u16);

impl LocalRef {
    pub fn inner(&self) -> u16 {
        self.0
    }
}

impl From<u16> for LocalRef {
    fn from(value: u16) -> Self {
        Self(value)
    }
}

/// Type context local to a specific function.
pub trait FnContext {
    /// Returns if the function is in a container.
    fn is_container(&self) -> bool;

    /// Returns the number of arguments.
    fn num_args(&self) -> u16;

    /// Returns the type of a function argument, if it exists.
    fn get_arg_type(&self, arg: ArgRef) -> Option<ValueType>;

    /// Returns the type of a local variable, if it exists.
    fn get_local_type(&self, local: LocalRef) -> Option<ValueType>;

    /// Returns the number of fields.
    fn num_instance_fields(&self) -> u16;

    /// Returns the type of the conatiner's field, if we're in a container and
    /// it exists.
    fn get_instance_field_type(&self, field: FieldRef) -> Option<ValueType>;
}

/// Facade trait to provide extra functionality to [`FnContext`].
pub trait FnContextExt {
    /// Resolves a SSA expr argument to a type, if possible.
    fn resolve_ssa_arg_type(&self, arg: &SsaArg) -> Option<ValueType>;

    /// Resolves an env expression to a type, if possible.
    fn resolve_env_expr_type(&self, env: &EnvExpr) -> Option<ValueType>;
}

impl<T: FnContext> FnContextExt for T {
    fn resolve_ssa_arg_type(&self, arg: &SsaArg) -> Option<ValueType> {
        match arg {
            SsaArg::Literal(lv) => Some(lv.ty()),
            SsaArg::Local(var) => self.get_local_type(*var),
            SsaArg::Arg(arg) => self.get_arg_type(*arg),
            SsaArg::Field(field) => {
                if self.is_container() {
                    self.get_instance_field_type(*field)
                } else {
                    None
                }
            }
        }
    }

    fn resolve_env_expr_type(&self, env: &EnvExpr) -> Option<ValueType> {
        match env {
            EnvExpr::Local(var) => self.get_local_type(*var),
            EnvExpr::Arg(arg) => self.get_arg_type(*arg),
            EnvExpr::Field(field) => {
                if self.is_container() {
                    self.get_instance_field_type(*field)
                } else {
                    None
                }
            }
        }
    }
}

/// Data about a container's composition.
///
/// For now, this is just a list of the field types.
pub struct ContainerData {
    field_types: Vec<ValueType>,
}

impl ContainerData {
    pub fn new(field_types: Vec<ValueType>) -> Self {
        Self { field_types }
    }
}

/// Data about a container's method.
///
/// For now, this is argument types and local variable types.
pub struct MethodData<'c> {
    container_data: &'c ContainerData,
    arg_types: Vec<ValueType>,
    local_types: Vec<ValueType>,
}

impl<'c> MethodData<'c> {
    pub fn new(
        container_data: &'c ContainerData,
        arg_types: Vec<ValueType>,
        local_types: Vec<ValueType>,
    ) -> Self {
        Self {
            container_data,
            arg_types,
            local_types,
        }
    }
}

impl<'c> FnContext for MethodData<'c> {
    fn is_container(&self) -> bool {
        true
    }

    fn num_args(&self) -> u16 {
        self.arg_types.len() as u16
    }

    fn get_arg_type(&self, arg: ArgRef) -> Option<ValueType> {
        self.arg_types.get(arg.inner() as usize).cloned()
    }

    fn get_local_type(&self, local: LocalRef) -> Option<ValueType> {
        self.local_types.get(local.inner() as usize).cloned()
    }

    fn num_instance_fields(&self) -> u16 {
        self.container_data.field_types.len() as u16
    }

    fn get_instance_field_type(&self, field: FieldRef) -> Option<ValueType> {
        self.container_data
            .field_types
            .get(field.inner() as usize)
            .cloned()
    }
}
